import numpy as np
import random

class RandomFlip:
    def __init__(self, h_prob=0.5, v_prob=0.5):
        self.h_prob = h_prob
        self.v_prob = v_prob
    
    def __call__(self, image, annotations, params=None):
        h, w = image.shape[:2]
        
        if random.random() < self.h_prob:
            image = image[:, ::-1]
            for i in range(len(annotations)):
                annotations[i]['bbox'][0], annotations[i]['bbox'][2] = w - annotations[i]['bbox'][2], w - annotations[i]['bbox'][0]

        if random.random() < self.v_prob:
            image = image[::-1, :]
            for i in range(len(annotations)):
                annotations[i]['bbox'][1], annotations[i]['bbox'][3] = h - annotations[i]['bbox'][3], h - annotations[i]['bbox'][1]

        return image, annotations

class RandomApply:
    def __init__(self, transform, p=0.5):
        self.transform = transform
        self.p = p
    
    def __call__(self, image, annotations, params=None):
        if random.random() < self.p:
            return self.transform(image, annotations, params=params)
        return image, annotations

class RandomCrop:
    def __init__(self, height_ratio=0.5, width_ratio=0.5):
        self.height_ratio = height_ratio
        self.width_ratio = width_ratio

    def __call__(self, image, annotations, params=None):
        h, w = image.shape[:2]

        new_h = int(h * self.height_ratio)
        new_w = int(w * self.width_ratio)

        x = int(random.random() * (w - new_w))
        y = int(random.random() * (h - new_h))

        image = image[y:y+new_h, x:x+new_w]

        clamp = lambda x, lower, upper: min(max(x, lower), upper)

        new_annotations = []
        for obj in annotations:
            inter_w = max(min(x+new_w, obj['bbox'][2]) - max(x, obj['bbox'][0]), 0)
            inter_h = max(min(y+new_h, obj['bbox'][3]) - max(y, obj['bbox'][1]), 0)
            inter_area_ratio = inter_h * inter_w / ((obj['bbox'][3] - obj['bbox'][1]) * (obj['bbox'][2] - obj['bbox'][0]))
            if inter_area_ratio < 0.5:
                continue

            obj['bbox'][0] = clamp(obj['bbox'][0] - x, 0, new_w)
            obj['bbox'][1] = clamp(obj['bbox'][1] - y, 0, new_h)
            obj['bbox'][2] = clamp(obj['bbox'][2] - x, 0, new_w)
            obj['bbox'][3] = clamp(obj['bbox'][3] - y, 0, new_h)
            obj['bbox'] = [float(i) for i in obj['bbox']]
            new_annotations.append(obj)

        return image, new_annotations


class FlexibleRandomCrop:
    def __init__(self, min_height_ratio=0.5, min_width_ratio=0.5):
        self.min_height_ratio = min_height_ratio
        self.min_width_ratio = min_width_ratio

    def __call__(self, image, annotations, params=None):
        h, w = image.shape[:2]

        height_ratio = random.random() * (1.0 - self.min_height_ratio) + self.min_height_ratio
        width_ratio = random.random() * (1.0 - self.min_width_ratio) + self.min_width_ratio

        new_h = int(h * height_ratio)
        new_w = int(w * width_ratio)

        x = int(random.random() * (w - new_w))
        y = int(random.random() * (h - new_h))

        image = image[y:y+new_h, x:x+new_w]

        clamp = lambda x, lower, upper: min(max(x, lower), upper)

        new_annotations = []
        for obj in annotations:
            inter_w = max(min(x+new_w, obj['bbox'][2]) - max(x, obj['bbox'][0]), 0)
            inter_h = max(min(y+new_h, obj['bbox'][3]) - max(y, obj['bbox'][1]), 0)
            inter_area_ratio = inter_h * inter_w / ((obj['bbox'][3] - obj['bbox'][1]) * (obj['bbox'][2] - obj['bbox'][0]))
            if inter_area_ratio < 0.5:
                continue

            obj['bbox'][0] = clamp(obj['bbox'][0] - x, 0, new_w)
            obj['bbox'][1] = clamp(obj['bbox'][1] - y, 0, new_h)
            obj['bbox'][2] = clamp(obj['bbox'][2] - x, 0, new_w)
            obj['bbox'][3] = clamp(obj['bbox'][3] - y, 0, new_h)
            obj['bbox'] = [float(i) for i in obj['bbox']]
            new_annotations.append(obj)

        return image, new_annotations